# Copyright (c) 2011, Chandler Armstrong (omni dot armstrong at gmail dot com)
# see LICENSE.txt for details




"""the frustum is world space"""




from __future__ import division

from pygame import image, Rect, sprite, Surface, transform
import os




N_HEADERLINES = 3 # the node data files use the first three lines for header data
MSG_ERRORSCALAR = """/
cull_scalar and frustum_scalar must be equal to or greater than 1/
"""



class _Frustum(object):


    def __init__(self, rect_viewer, pos, node, graph,
                 cull_scalar=1, frustum_scalar=1):
        # construct frustum object
        # size -- the size of the frustum
        # pos -- the initial position in a node space
        # node -- the initial node
        # graph -- the graph object specificing node topology
        # cull_scalar -- the scalar to increase cull frustum size over frustum
        #   size
        if cull_scalar < 1 or frustum_scalar < 1:
            raise ValueError, MSG_ERRORSCALAR
        get_data, get_img = self._get_data, self._get_img

        # center frustum rect on viewer
        offset = [frustum_scalar * d for d in rect_viewer.size]
        rect = rect_viewer.inflate(*offset)
        rect_cull = rect.inflate(rect.w * cull_scalar, rect.h * cull_scalar)

        # initialize nodes; one node must be in nodes to initialize culling
        graph_nodes = dict(graph.nodes(data=True))
        graph_node = graph_nodes[node]
        node_rect, node_data = get_data(graph_node['data'])
        node_img = get_img(graph_node['img'])
        # move node so that center of frustum rect is on node rect at pos
        node_rect.topleft = rect.center
        x, y = -pos[0], -pos[1]
        node_rect.move_ip((x, y))
        # initialize nodes with the first node
        nodes = {node : {'rect' : node_rect,
                         'data' : node_data,
                         'img' : node_img}}

        # current node and position on node space of frustum center
        self.pos = pos
        self.node = node
        # loading and culling rects
        self.rect = rect
        self._rect_cull = rect_cull
        # graph and loaded nodes
        self._graph = graph
        self.nodes = nodes
        # surface
        self.surf = Surface(rect.size)

        self._cull()
        self.draw()


    def __getstate__(self):
        odict = self.__dict__.copy()
        #for node in odict.nodes.values(): node.clear()
        del odict['nodes']
        del odict['surf']
        return odict


    def __setstate__(self, ndict):
        self.__dict__.update(ndict)
        get_data, get_img = self._get_data, self._get_img
        graph, node, pos, rect = self._graph, self.node, self.pos, self.rect

        # initialize nodes; one node must be in nodes to initialize culling
        graph_nodes = dict(graph.nodes(data=True))
        graph_node = graph_nodes[node]
        node_rect, node_data = get_data(graph_node['data'])
        node_img = get_img(graph_node['img'])
        # move node so that center of frustum rect is on node rect at pos
        node_rect.topleft = rect.center
        x, y = -pos[0], -pos[1]
        node_rect.move_ip((x, y))
        # initialize nodes with the first node
        nodes = {node : {'rect' : node_rect,
                         'data' : node_data,
                         'img' : node_img}}
        self.nodes = nodes
        self.surf = Surface(rect.size)
        self._cull()
        self.draw()


    def _get_data(self, p_data):
        # get data and rect from data file
        # p_data -- the path to node data file
        f_data = open(os.path.realpath(p_data))

        header = dict(f_data.readline().split() for n in xrange(N_HEADERLINES))
        data = f_data.readlines()

        ncols, nrows, cellsize = (int(header['ncols']),
                                  int(header['nrows']),
                                  int(header['cellsize']))

        size = (ncols * cellsize, nrows * cellsize)
        rect = Rect((0, 0), size)

        return rect, data


    def _get_img(self, p_img):
        # get img from img file
        # p_img -- the path to node img file
        surf = image.load(os.path.realpath(p_img))
        surf.convert()
        return surf


    def _get_node(self, node, topleft=None):
        # get node and position in world space
        # node -- the node to get
        # topleft -- the position of node in world space
        graph_nodes = dict(self._graph.nodes(data=True))
        graph_node = graph_nodes[node]
        node_rect, node_data = self._get_data(graph_node['data'])
        node_img = self._get_img(graph_node['img'])
        if topleft: node_rect.topleft = topleft
        return {'rect' : node_rect, 'data' : node_data, 'img' : node_img}


    def _cull(self):
        # add nodes that have entered frustum, remove nodes that have exited frustum
        # cull will not remove self.node from nodes
        # cull must be called prior to updating self.node during a move operation
        rect, rect_cull = self.rect, self._rect_cull
        nodes, graph = self.nodes, self._graph
        get_node = self._get_node

        que = [self.node]
        checked = set()
        while que:
            node = que.pop()
            if node in checked: continue
            checked.add(node)
            node_rect = nodes[node]['rect']
            # check evict node
            if not rect_cull.colliderect(node_rect):
                if node == self.node: continue # do not remove current node
                del nodes[node]
                continue # skip adding neighbors
            # check add neighbors
            for neighbor, atts in graph[node].items():
                if not neighbor in nodes:
                    offset = (atts['x'], atts['y'])
                    topleft = node_rect.move(offset).topleft
                    nodes[neighbor] = get_node(neighbor, topleft=topleft)

                que = [neighbor] + que
        # if a node was not checked, it cannot be in the frustum
        for node in set(nodes).difference(checked): del nodes[node]


    def worldpos2nodepos(self, pos, node=None):
        # convert worldpos to nodepos
        # pos -- the worldpos
        # node -- the initial node in the search for nodepos
        #   NOTE: node must be in nodes; cull guarantees self.node is in nodes
        graph, nodes = self._graph, self.nodes.copy()
        get_node = self._get_node
        if not node: node = nodes.keys()[0]

        # main
        que = [node]
        checked = set()
        while que:
            node = que.pop()
            if node in checked: continue
            checked.add(node)
            node_rect = nodes[node]['rect']
            # check if pos maps to this node
            if node_rect.collidepoint(pos):
                que = True # flag that a node was found
                break
            # check add neighbors
            for neighbor, atts in graph[node].items():
                if not neighbor in nodes:
                    offset = (atts['x'], atts['y'])
                    topleft = node_rect.move(offset).topleft
                    nodes[neighbor] = get_node(neighbor, topleft=topleft)

                que = [neighbor] + que

        # check that a node was found
        if not que:
            raise ValueError(
                'worldpos (%i, %i) does not map to any node' % (pos[0], pos[1]))
        # return position on node and node id
        return (pos[0] - node_rect.x, pos[1] - node_rect.y), node


    def nodepos2worldpos(self, (x, y), node):
        # convert nodepos to worldpos
        nodes = self.nodes
        rect = nodes[node]['rect']
        return x + rect.x, y + rect.y


    def move(self, (x, y)):
        # move node rects on frustum
        # (x, y) -- the movement vector
        node, nodes = self.node, self.nodes

        worldpos = self.nodepos2worldpos(self.pos, node)
        # move node rects on frustum
        for v in nodes.values(): v['rect'].move_ip((-x, -y))
        self._cull() # cull must be called before setting new nodepos

        self.pos, self.node = self.worldpos2nodepos(worldpos, node)


    def draw(self):
        # draw nodes to frustum surface
        nodes, surf = self.nodes, self.surf

        for v in nodes.values():
            img = v['img']
            rect = v['rect']
            surf.blit(img, rect)



class Main(object):
    """
    Interface with frustum object

    attributes:
    rect -- rect object to position viewer on the scren
    area -- the area of the visible frustum image to be blitted.  area is the
      value for the 'area' keyword argument of the Surface.blit method.  for
      example:

      SCREEN.blit(viewer.surf, viewer.rect, area=viewer.area)

      the area value is automatically scaled for zoom level.  the area attribute
      is read only.

    surf -- the surf attribute is what should be blitted to the SCREEN or other
      surface.  the surf is be the current frustum image with all active sprites
      drawn.  the surf is scaled for the zoom level.  the surf is a read only.
    group_evict -- the evicted, or inactive, sprite group.  the viewer will
      check this group for sprites that have entered the frustum every time
      update0 is called.  the viewer will place sprites that have entered the
      frustum into an active sprite group that is updated by the viewer's
      update0 method.  sprites in the active group will be automatically
      transformed and repositioned as the viewer is scrolled and zoomed.  place
      all sprites initially in the 'group_evict' attribute; the inactive sprite
      group. the viewer will automatically handle sorting them.
    group_active -- the active sprite group.  do not place sprites directly in
      the active sprite group.  always place sprites in the inactive sprite
      group.  the viewer will automatically handle sorting them.  only access
      the active sprite group to update sprites or iterate through sprites to
      test for collisions.

    NOTE: the update method occurs in two steps, update0 and update1.  update0
    updates the world position of each sprite, and update1 updates the node
    position of each sprite from the updated world position.  updating world
    position and node position separately permit client code to operate upon
    sprite world positions following update0, before the sprites new node
    position is set by update1.  for example, after updating world positions
    with update0, client code could respond to collisions, and adjust sprite
    world positions if necessary, before updating sprite node positions.
    """


    def __init__(self, pos, node, graph, rect,
                 zoom_in_limit=2, zoom_out_limit=0.001, **kw):
        """
        arguments:
        pos -- position on initial node rect
        node -- initial node key
        graph -- graph object encoding node topology
        rect -- rect object representing viewer size and position on the screen

        keyword arguments:
        zoom_in_limit -- viewer zoom in limit
        zoom_out_limit -- viewer zoom out limit
        cull_scalar -- scalar of cull size relative to frustum size;
          must be <= 1
        frustum_scalar -- scalar of frustum size relative to viewer size;
          must be <= 1
        """
        self.frustum = _Frustum(rect, pos, node, graph, **kw)
        self.rect = rect
        self.group_evict = sprite.Group()
        self.group_active = sprite.Group()
        self._area = Rect((0, 0), rect.size)
        self._scalar = 1 # zoom size scalar
        self._surf = self.frustum.surf
        self._zoom_in_limit = zoom_in_limit
        self._zoom_out_limit = zoom_out_limit


    def __getstate__(self):
        odict = self.__dict__.copy()
        del odict['_surf']
        return odict


    def __setstate__(self, ndict):
        self.__dict__.update(ndict)
        self._surf = self.frustum.surf


    @property
    def area(self):
        """the area scaled and positioned for zoom level"""
        area, scalar = self._area, self._scalar
        rect = area.copy()
        x, y = rect.center
        rect.center = (x * scalar, y * scalar)
        return rect


    @property
    def surf(self):
        """
        the frustum surface with all sprites drawn.  the surf is scaled for the
        zoom level
        """
        scalar, surf = self._scalar, self._surf
        w, h = surf.get_size()
        size = (w * scalar, h * scalar)
        return transform.smoothscale(surf, size)


    def scroll(self, x, y):
        """
        Scroll viewer frustum area.

        arguments:
        x -- x scroll amount
        y -- y scroll amount

        return None
        """
        # cache objects
        frustum, scalar = self.frustum, self._scalar
        group_active, group_evict = self.group_active, self.group_evict
        nodes = frustum.nodes
        # move frustum
        frustum.move((x * scalar, y * scalar))
        frustum.draw()
        #self._surf = frustum.surf
        # check if evicted sprites have entered frustum
        for spr in group_evict:
            if spr.node in nodes:
                group_evict.remove(spr)
                group_active.add(spr)


    def zoom(self, z):
        """
        Zoom viewer frustum area

        arguments:
        z -- zoom amount, positive zooms in, negative zooms out

        return None
        """
        scalar = self._scalar + z
        if self._zoom_in_limit >= scalar >= self._zoom_out_limit:
            self._scalar = scalar


    def update0(self):
        """
        sort sprites into the active and inactive sprite groups, depending upon
        if the node they are upon is currently within the frustum.  get and set
        the world position for each sprite.  after calling update0, client code
        should do any further processing, including sprite updates and
        collision handling, using the newly supplied world position of the
        sprites stored in the sprite rect attribute, before calling update1.

        return None

        NOTE: scroll must be called prior to update, to ensure that frustum
        nodes are current
        """
        # cache objects
        frustum = self.frustum
        group_active, group_evict = self.group_active, self.group_evict
        nodes, nodepos2worldpos = frustum.nodes, frustum.nodepos2worldpos
        # update spr.rect to the worldpos
        for spr in group_active.sprites():
            # check that active sprite still active
            if not spr.node in nodes:
                group_active.remove(spr)
                group_evict.add(spr)
                continue
            # reset to a fresh rect from vzrect each loop, draw will scale the rect
            vzrect = spr.vzrect
            rect = Rect((0, 0), vzrect.size)
            rect.center = nodepos2worldpos(vzrect.center, spr.node)
            spr.P = [nodepos2worldpos(*vzp) for vzp in spr.vzP]
            spr.rect = rect


    def update1(self):
        """
        get and set the node position of each sprite from its world position.
        check if any sprite has exited the culling frustum and remove that
        sprite from the active sprite group and add it to the inactive sprite
        group if the sprite has exited the culling frustum

        return None

        NOTE: update0 must be called prior to update1
        """
        group_active = self.group_active
        worldpos2nodepos = self.frustum.worldpos2nodepos
        for spr in group_active.sprites():
            vzpos = worldpos2nodepos(spr.rect.center, spr.node)
            spr.vzP = [worldpos2nodepos(p) for p in spr.P]
            spr.vzrect.center, spr.node = vzpos


    def draw(self):
        """
        draw sprite group to surface
        draw will scale sprite rects appropriately prior to drawing

        return a copy of the prior surface, used for the clear method

        NOTE: update0 and update1 must be called prior to draw
        """
        group, surf = self.group_active, self._surf
        surface_clear = surf.copy() # a surface to clear drawn sprites
        group.draw(surf)
        return surface_clear


    def clear(self, surface_clear):
        """
        Clear surface

        arguments:
        surface_clear -- the surface used to erase old sprite positions

        return None

        NOTE: draw must be called prior to clear, and the surface returned by
          draw provided to clear
        """
        self.group_active.clear(self._surf, surface_clear)


    def get_mouse(self, (x, y), nodespace=False):
        """
        Return mouse position in either world or node space

        if the mouse is not on the viewer, an error will be raised

        returns an x, y coordinates on viewer rect
        in nodespace == True
          returns an x, y value and the node id
          if the mouse is on a node:
            returns x, y coordinates on node and that node's id
          if the mouse is not on a node:
            returns the x, y coordinates on the viewer and a None value for the
              node id
        """
        nodes, rect = self.frustum.nodes, self.rect
        # convert (x, y) to position on viewer rect
        x -= rect.x
        y -= rect.y
        if nodespace == False: return x, y
        # get x, y in nodespace
        if 0 > x > (rect.x + rect.w) or 0 > y > (rect.y + rect.h):
            raise ValueError, "mouse not on viewer"

        for node in nodes:
            node_rect = nodes[node]['rect']
            if node_rect.collidepoint((x, y)):
                # convert (x, y) to position on node rect and node id
                x -= node_rect.x
                y -= node_rect.y
                return (x, y), node
        # return position on viewer rect and None for node id
        return (x, y), None
